import asyncio
import json
import os
import sys
from typing import Optional, Dict, Any
from contextlib import AsyncExitStack
import logging
import time

from mcp import ClientSession
from mcp.client.sse import sse_client

import dashscope
from dashscope import Generation
from dotenv import load_dotenv

# 配置日志
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

load_dotenv()  # load environment variables from .env

class MCPClient:
    def __init__(self):
        # Initialize session and client objects
        self.session: Optional[ClientSession] = None
        self._session_context = None
        self._streams_context = None
        
        # 初始化通义千问客户端
        self.api_key = os.getenv('DASHSCOPE_API_KEY')
        if not self.api_key:
            raise ValueError("DASHSCOPE_API_KEY not found in environment variables")
        dashscope.api_key = self.api_key
        logger.info("Initialized with DashScope API key")
        
        # 是否处于调试模式
        self.debug = True
        
        # 接收到的SSE事件计数
        self.event_count = 0

    async def connect_to_sse_server(self, server_url: str):
        """Connect to an MCP server running with SSE transport"""
        logger.info(f"Connecting to SSE server at {server_url}")
        
        # 确保 server_url 是正确的 SSE 端点
        if not server_url.endswith('/sse'):
            server_url = f"{server_url.rstrip('/')}/sse"
            logger.info(f"Adjusted SSE URL to {server_url}")
            
        # Store the context managers so they stay alive
        try:
            # 直接测试SSE连接
            logger.info(f"Testing SSE connection directly...")
            
            self._streams_context = sse_client(url=server_url)
            streams = await self._streams_context.__aenter__()
            logger.info(f"SSE stream established")
            
            # 监听SSE事件
            async def handle_sse_event(event):
                self.event_count += 1
                event_data = event.data if hasattr(event, 'data') else None
                event_type = event.type if hasattr(event, 'type') else 'message'
                
                if not event_data:
                    logger.debug(f"Received SSE event #{self.event_count} of type {event_type} without data")
                    return
                    
                try:
                    data = json.loads(event_data)
                except json.JSONDecodeError:
                    logger.warning(f"Received SSE event #{self.event_count} with invalid JSON data: {event_data}")
                    return
                    
                if not isinstance(data, dict):
                    logger.warning(f"Received SSE event #{self.event_count} with non-dict data: {data}")
                    return
                    
                method = data.get('method')
                params = data.get('params', {})
                
                if method == 'heartbeat':
                    count = params.get('count', 'N/A')
                    logger.info(f"Received heartbeat #{count}")
                    print(f"Heartbeat received: #{count}")
                    return
                    
                if not method or not params:
                    logger.warning(f"Received invalid event format: {data}")
                    return
                    
                logger.info(f"Processing event: method={method}, params={params}")
                
                query_content = params.get('content', '')
                if not query_content:
                    logger.warning("Received event without query content")
                    return
                    
                try:
                    result = await self.process_query(query_content)
                    logger.info(f"Processed event result: {result}")
                    
                    # 将处理结果发送回服务器
                    response = await self.session.notify("response", {"content": result})
                    logger.info(f"Response sent to server: {response}")
                    print(f"\nResponse: {result}")
                except Exception as e:
                    error_msg = f"Error processing query: {str(e)}"
                    logger.error(error_msg)
                    print(f"\nError: {error_msg}")
                    # 发送错误响应给服务器
                    try:
                        await self.session.notify("error", {"message": error_msg})
                    except Exception as e2:
                        logger.error(f"Failed to send error response: {e2}")
            
            # 添加事件处理器
            streams[0].on_message = handle_sse_event
            streams[0].on_error = lambda e: logger.error(f"SSE Error: {str(e)}")
            streams[0].on_open = lambda: logger.info("SSE connection opened")
            streams[0].on_close = lambda: logger.info("SSE connection closed")

            self._session_context = ClientSession(*streams)
            self.session = await self._session_context.__aenter__()
            logger.info(f"MCP session initialized")

            # Initialize
            logger.info(f"Initializing MCP session")
            await self.session.initialize()
            logger.info(f"MCP session initialized successfully")

            # 验证连接并获取工具列表
            print("Initialized SSE client...")
            print("Listing tools...")
            
            max_retries = 3
            retry_delay = 2
            
            for attempt in range(max_retries):
                try:
                    logger.info(f"Requesting tool list (attempt {attempt + 1}/{max_retries})")
                    response = await asyncio.wait_for(self.session.list_tools(), timeout=10.0)
                    tools = response.tools
                    logger.info(f"Received {len(tools)} tools from server")
                    print("\nConnected to server with tools:", [tool.name for tool in tools])
                    break
                except asyncio.TimeoutError:
                    logger.warning(f"Timeout on attempt {attempt + 1}/{max_retries}")
                    if attempt == max_retries - 1:
                        logger.error("All retries exhausted")
                        if self.event_count == 0:
                            raise TimeoutError("服务器连接超时：未收到任何SSE事件，请检查服务器状态")
                        else:
                            raise TimeoutError(f"服务器响应超时：已收到{self.event_count}个SSE事件，但工具列表请求失败")
                    await asyncio.sleep(retry_delay)
            
        except Exception as e:
            logger.error(f"Error connecting to SSE server: {str(e)}")
            raise

    async def cleanup(self):
        """Properly clean up the session and streams"""
        logger.info("Cleaning up resources")
        if self._session_context:
            await self._session_context.__aexit__(None, None, None)
        if self._streams_context:
            await self._streams_context.__aexit__(None, None, None)

    async def debug_call_tool(self, tool_name: str, parameters: Dict[str, Any]) -> Dict:
        """直接测试工具调用，不通过大模型"""
        if not self.session:
            raise ValueError("Session not initialized")
            
        logger.info(f"Directly calling tool: {tool_name} with parameters: {parameters}")
        start_time = time.time()
        
        try:
            result = await asyncio.wait_for(
                self.session.call_tool(tool_name, parameters), 
                timeout=10.0
            )
            logger.info(f"Tool result received in {time.time() - start_time:.2f}s")
            return result
        except asyncio.TimeoutError:
            logger.error(f"Tool call timed out after 10 seconds")
            raise TimeoutError(f"Tool call to {tool_name} timed out")
        except Exception as e:
            logger.error(f"Tool call error: {str(e)}")
            raise

    async def process_query(self, query: str) -> str:
        """Process a query using DashScope API and available tools"""
        logger.info(f"Processing query: {query}")
        
        # 添加调试快捷方式
        if query.startswith("!debug "):
            parts = query[7:].split(" ", 1)
            if len(parts) == 2 and parts[0] == "weather":
                try:
                    logger.info(f"Debug mode: directly calling get_weather with city={parts[1]}")
                    result = await self.debug_call_tool("get_weather", {"city": parts[1]})
                    return f"Debug tool call result: {result.content}"
                except Exception as e:
                    return f"Debug error: {str(e)}"
        
        if query.startswith("!test"):
            try:
                logger.info("Running basic tool test")
                result = await self.debug_call_tool("hello_world", {"name": "Debug Test"})
                return f"Test result: {result.content}"
            except Exception as e:
                return f"Test error: {str(e)}"
                
        messages = [
            {
                "role": "user",
                "content": query
            }
        ]

        try:
            logger.info("Requesting tool list for API call")
            tools_response = await asyncio.wait_for(
                self.session.list_tools(), 
                timeout=5.0
            )
            
            available_tools = []
            for tool in tools_response.tools:
                logger.info(f"Processing tool: {tool.name}")
                tool_info = {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": {
                        "type": "object",
                        "properties": {},
                        "required": []
                    }
                }
                
                if hasattr(tool, 'inputSchema') and tool.inputSchema:
                    try:
                        if isinstance(tool.inputSchema, str):
                            schema = json.loads(tool.inputSchema)
                        else:
                            schema = tool.inputSchema
                        tool_info["parameters"] = schema
                    except Exception as e:
                        logger.warning(f"Failed to parse input schema for tool {tool.name}: {e}")
                
                available_tools.append(tool_info)
            
            logger.info(f"Available tools for API call: {[tool['name'] for tool in available_tools]}")

            # Initial DashScope API call
            logger.info("Making initial API call to DashScope")
            start_time = time.time()
            response = Generation.call(
                model='qwen-max',
                messages=messages,
                tools=available_tools,
                tool_choice="auto",
                result_format='message'
            )
            logger.info(f"DashScope API call completed in {time.time() - start_time:.2f}s")

            # Process response and handle tool calls
            tool_results = []
            final_text = []

            if response.status_code == 200:
                logger.info("Received successful response from DashScope")
                output = response.output
                logger.info(f"Response output: {output}")
                assistant_message = output.choices[0].message
                
                logger.info(f"Assistant message: {assistant_message}")
                
                if hasattr(assistant_message, 'content') and assistant_message.content:
                    final_text.append(assistant_message.content)
                
                # 检查是否有函数调用
                if hasattr(assistant_message, 'function_call') and assistant_message.function_call:
                    tool_name = assistant_message.function_call.name
                    try:
                        tool_args = json.loads(assistant_message.function_call.arguments)
                    except:
                        tool_args = assistant_message.function_call.arguments
                    
                    logger.info(f"Calling tool: {tool_name} with args: {tool_args}")
                    
                    # Execute tool call
                    try:
                        tool_start = time.time()
                        result = await asyncio.wait_for(
                            self.session.call_tool(tool_name, tool_args),
                            timeout=10.0
                        )
                        logger.info(f"Tool call completed in {time.time() - tool_start:.2f}s")
                        tool_results.append({"call": tool_name, "result": result})
                        logger.info(f"Tool result: {result.content}")
                        final_text.append(f"[Tool {tool_name} returned: {result.content}]")

                        # Continue conversation with tool results
                        messages.extend([
                            {
                                "role": "assistant",
                                "content": assistant_message.content if hasattr(assistant_message, 'content') else None,
                                "function_call": {
                                    "name": tool_name,
                                    "arguments": json.dumps(tool_args)
                                }
                            },
                            {
                                "role": "function",
                                "name": tool_name,
                                "content": result.content
                            }
                        ])

                        # Get next response from DashScope
                        logger.info("Making follow-up API call to DashScope")
                        follow_up_start = time.time()
                        response = Generation.call(
                            model='qwen-max',
                            messages=messages,
                            tools=available_tools,
                            tool_choice="auto",
                            result_format='message'
                        )
                        logger.info(f"Follow-up API call completed in {time.time() - follow_up_start:.2f}s")

                        if response.status_code == 200:
                            assistant_message = response.output.choices[0].message
                            if hasattr(assistant_message, 'content') and assistant_message.content:
                                final_text.append(assistant_message.content)
                    except asyncio.TimeoutError:
                        error_msg = f"Tool call to {tool_name} timed out after 10 seconds"
                        logger.error(error_msg)
                        final_text.append(f"Error: {error_msg}")
                    except Exception as e:
                        error_msg = f"Error calling tool {tool_name}: {str(e)}"
                        logger.error(error_msg)
                        final_text.append(f"Error: {error_msg}")
            else:
                error_msg = f"API call failed with status code: {response.status_code}"
                logger.error(error_msg)
                raise Exception(error_msg)

            result = "\n".join([t for t in final_text if t])
            logger.info(f"Final result: {result}")
            return result
            
        except asyncio.TimeoutError:
            error_msg = "Operation timed out"
            logger.error(error_msg)
            return f"Error: {error_msg}"
        except Exception as e:
            error_msg = f"Error processing query: {str(e)}"
            logger.error(error_msg)
            return f"Error: {error_msg}"

    async def chat_loop(self):
        """Run an interactive chat loop"""
        print("\nMCP Client Started!")
        print("Type your queries or 'quit' to exit.")
        print("Special commands:")
        print("  !debug weather <city> - Test weather tool directly")
        print("  !test - Test basic tool functionality")
        
        while True:
            try:
                query = input("\nQuery: ").strip()
                
                if query.lower() == 'quit':
                    logger.info("User requested to quit")
                    break
                    
                start_time = time.time()
                print("Processing query...")
                response = await self.process_query(query)
                logger.info(f"Query processed in {time.time() - start_time:.2f}s")
                print("\nResponse:", response)
                    
            except Exception as e:
                error_msg = f"Error in chat loop: {str(e)}"
                logger.error(error_msg)
                print(f"\nError: {str(e)}")

async def main():
    if len(sys.argv) < 2:
        print("Usage: python client_sse.py <URL of SSE MCP server (i.e. http://localhost:8000/sse)>")
        sys.exit(1)
        
    server_url = sys.argv[1]
    client = MCPClient()
    try:
        await client.connect_to_sse_server(server_url=server_url)
        await client.chat_loop()
    except Exception as e:
        print(f"Error: {str(e)}")
        print("Make sure the server is running and the URL is correct.")
        print("Example usage: python client_sse.py http://localhost:8000")
    finally:
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())